#include "SimpleStrokeOptimizer.h"
#include "Scene.h"
#include "FluidSparseMatrix.h"
#include "CSRMatrix.h"
#include "MatrixSolver.h"
#include "PCGSolver.h"
#include "GridField.h"
#include "mathlib.h"
#include "DecoArticulatedBody.h"
#include "DecoLogger.h"

SimpleStrokeOptimizer::SimpleStrokeOptimizer(DecoScene* scene) : mB(NULL), mF(NULL), mTmp(NULL), mRhs(NULL), mGrid(NULL), mCSRMat(NULL), mFluidMat(NULL), mScene(scene)
{

}

SimpleStrokeOptimizer::~SimpleStrokeOptimizer()
{
	clear();
}

void SimpleStrokeOptimizer::SetFluidField(GridField* field)
{
	mGrid = field;
}

void SimpleStrokeOptimizer::SetWeights(const std::vector<double>& weights)
{
	mWeights = weights;
}
void SimpleStrokeOptimizer::SetDesiredDirection(const vector3& dir)
{
	mDesiredDir = dir;
}
double SimpleStrokeOptimizer::GetOptimalQDot(double timeStep, std::vector<double>& qdot)
{
	return optimize(timeStep, qdot);
}

void SimpleStrokeOptimizer::SetMatrices(CSRMatrix* csrMat, FLuidSparseMatrix* fluidMat)
{
	if (!csrMat || !fluidMat)
	{
		(*DecoLogger::GetSingleton()) << "Set NULL matrices in SimpleStrokeOptimizer, check whether bAccurateCoupling is on.\n";
		return;
	}
	mCSRMat = csrMat;
	mFluidMat = fluidMat;
}


double SimpleStrokeOptimizer::optimize(double timeStep, std::vector<double>& qdot)
{
	calculateBMatrix(timeStep);
	calculateFMatrix(timeStep);
	PCGSolver solver;
	solver.SetMatrix(mCSRMat);
	int nDofs = static_cast<int>(mWeights.size());
	FLuidSparseMatrix* fluidMat = mFluidMat;
	int dim = fluidMat->GetDim();
	for (int i = 0; i < nDofs; ++i)
	{
		if (mWeights[i] < EPSILON_FLOAT)
			continue;
		solver.SetRHS(mB[i]);
		solver.Solve(mTmp[i]);
	}
	assert(!mRhs);
	mRhs = new double[nDofs];
	for (int i = 0; i < nDofs; ++i)
	{
		mRhs[i] = 0;
		if (mWeights[i] < EPSILON_FLOAT)
			continue;
		for (int j = 0; j < 3; ++j)
		{
			double tmpSum[3];
			tmpSum[j] = 0;
			for (int k = 0; k < dim; ++k)
			{
				tmpSum[j] += mF[j][k] * mTmp[i][k];
			}
			mRhs[i] += mDesiredDir[j] * tmpSum[j];
		}
	}
	DecoArticulatedObject* artiObj = getArticulatedObject();
	artiObj->GetQDot(qdot);
	for (int i = 0; i < nDofs; ++i)
	{
		if (mWeights[i] < EPSILON_FLOAT)
		{			
			continue;
		}
		qdot[i] += mRhs[i] / mWeights[i] * timeStep;
	}
	return calculateOptimalValue(qdot);
}
void SimpleStrokeOptimizer::calculateFMatrix(double timeStep)
{
	FLuidSparseMatrix* fluidMat = mFluidMat;
	int dim = fluidMat->GetDim();
	int nDofs = static_cast<int>(mWeights.size());
	assert(!mF);
	mF = Malloc2DArray<double>(3, dim, TRUE);
	DecoArticulatedObject* artiObj = getArticulatedObject();
	double mass = artiObj->GetMass();
	for (int i = 0; i < dim; ++i)
	{
		Index3 index = fluidMat->GetIndex3D(i);
		vector3 force(0, 0, 0);
		double faceArea = mGrid->stepSize * mGrid->stepSize;
		if (mGrid->isNearSolid(index.m_i, index.m_j, index.m_k))
		{
			if (mGrid->isSolidCell(index.m_i + 1, index.m_j, index.m_k) && artiObj->isGridCellBelongsTo(index.m_i + 1, index.m_j, index.m_k))
			{
				force += faceArea / mass * vector3(1, 0, 0);
			}
			if (mGrid->isSolidCell(index.m_i - 1, index.m_j, index.m_k) && artiObj->isGridCellBelongsTo(index.m_i - 1, index.m_j, index.m_k))
			{
				force += faceArea / mass * vector3(-1, 0, 0);
			}
			if (mGrid->isSolidCell(index.m_i, index.m_j + 1, index.m_k) && artiObj->isGridCellBelongsTo(index.m_i, index.m_j + 1, index.m_k))
			{
				force += faceArea / mass * vector3(0, 1, 0);
			}
			if (mGrid->isSolidCell(index.m_i, index.m_j - 1, index.m_k) && artiObj->isGridCellBelongsTo(index.m_i, index.m_j - 1, index.m_k))
			{
				force += faceArea / mass * vector3(0, -1, 0);
			}
			if (mGrid->isSolidCell(index.m_i, index.m_j, index.m_k + 1) && artiObj->isGridCellBelongsTo(index.m_i, index.m_j, index.m_k + 1))
			{
				force += faceArea / mass * vector3(0, 0, 1);
			}
			if (mGrid->isSolidCell(index.m_i, index.m_j, index.m_k - 1) && artiObj->isGridCellBelongsTo(index.m_i, index.m_j, index.m_k - 1))
			{
				force += faceArea / mass * vector3(0, 0, -1);
			}
			mF[0][i] = force.x;
			mF[1][i] = force.y;
			mF[2][i] = force.z;
		}
	}
}

void SimpleStrokeOptimizer::calculateBMatrix(double timeStep)
{
	FLuidSparseMatrix* fluidMat = mFluidMat;
	int dim = fluidMat->GetDim();
	int nDofs = static_cast<int>(mWeights.size());
	assert(!mB);
	assert(!mTmp);
	mB = Malloc2DArray<double>(nDofs, dim, TRUE);
	mTmp = Malloc2DArray<double>(nDofs, dim, TRUE);
	DecoArticulatedObject* artiObj = getArticulatedObject();
	for (int i = 0; i < dim; ++i)
	{
		Index3 index = fluidMat->GetIndex3D(i);
		if (mGrid->isNearSolid(index.m_i, index.m_j, index.m_k))
		{
			if (mGrid->isSolidCell(index.m_i + 1, index.m_j, index.m_k) && artiObj->isGridCellBelongsTo(index.m_i + 1, index.m_j, index.m_k))
			{
				Matd jacobian(3, nDofs);
				bool succeed = artiObj->GetJacobianAtUFace(index.m_i + 1, index.m_j, index.m_k, jacobian);
				if (succeed)
					for (int j = 0; j < nDofs; ++j)
						mB[j][i] -= jacobian[0][j] / timeStep;
			}
			if (mGrid->isSolidCell(index.m_i - 1, index.m_j, index.m_k) && artiObj->isGridCellBelongsTo(index.m_i - 1, index.m_j, index.m_k))
			{
				Matd jacobian(3, nDofs);
				bool succeed = artiObj->GetJacobianAtUFace(index.m_i, index.m_j, index.m_k, jacobian);
				if (succeed)
					for (int j = 0; j < nDofs; ++j)
						mB[j][i] += jacobian[0][j] / timeStep;
			}
			if (mGrid->isSolidCell(index.m_i, index.m_j + 1, index.m_k) && artiObj->isGridCellBelongsTo(index.m_i, index.m_j + 1, index.m_k))
			{
				Matd jacobian(3, nDofs);
				bool succeed = artiObj->GetJacobianAtVFace(index.m_i, index.m_j + 1, index.m_k, jacobian);
				if (succeed)
					for (int j = 0; j < nDofs; ++j)
						mB[j][i] -= jacobian[1][j] / timeStep;
			}
			if (mGrid->isSolidCell(index.m_i, index.m_j - 1, index.m_k) && artiObj->isGridCellBelongsTo(index.m_i, index.m_j - 1, index.m_k))
			{
				Matd jacobian(3, nDofs);
				bool succeed = artiObj->GetJacobianAtVFace(index.m_i, index.m_j, index.m_k, jacobian);
				if (succeed)
					for (int j = 0; j < nDofs; ++j)
						mB[j][i] += jacobian[1][j] / timeStep;
			}
			if (mGrid->isSolidCell(index.m_i, index.m_j, index.m_k + 1) && artiObj->isGridCellBelongsTo(index.m_i, index.m_j, index.m_k + 1))
			{
				Matd jacobian(3, nDofs);
				bool succeed = artiObj->GetJacobianAtWFace(index.m_i, index.m_j, index.m_k + 1, jacobian);
				if (succeed)
					for (int j = 0; j < nDofs; ++j)
						mB[j][i] -= jacobian[2][j] / timeStep;
			}
			if (mGrid->isSolidCell(index.m_i, index.m_j, index.m_k - 1) && artiObj->isGridCellBelongsTo(index.m_i, index.m_j, index.m_k - 1))
			{
				Matd jacobian(3, nDofs);
				bool succeed = artiObj->GetJacobianAtWFace(index.m_i, index.m_j, index.m_k, jacobian);
				if (succeed)
					for (int j = 0; j < nDofs; ++j)
						mB[j][i] += jacobian[2][j] / timeStep;
			}
		}

	}
}

void SimpleStrokeOptimizer::clear()
{
	FLuidSparseMatrix* fluidMat = mFluidMat;
	int dim = fluidMat->GetDim();

	if (mB)
		Free2DArray<double>(mB, static_cast<UINT>(mWeights.size()), dim);
	if (mF)
		Free2DArray<double>(mF, 3, dim);
	if (mTmp)
		Free2DArray<double>(mTmp, static_cast<UINT>(mWeights.size()), dim);
	if (mRhs)
		delete[] mRhs;
}

DecoArticulatedObject* SimpleStrokeOptimizer::getArticulatedObject()
{
	return mScene->GetArticulatedObject();

}

double SimpleStrokeOptimizer::calculateOptimalValue(const std::vector<double>& qdot)
{
	double ret = 0;
	for (size_t i = 0; i < mWeights.size(); ++i)
	{
		ret += 0.5 * qdot[i] * qdot[i] * mWeights[i];
		ret -= mRhs[i] * qdot[i];
	}
	return ret;
}